package com.xiayuanxing.arithmetic.DataStructures.HuffmanCode;

import java.io.*;
import java.util.*;

/**
 * @program: arithmetic
 * @description: 赫夫曼编码
 * @author: xia yuan xing
 * @create: 2021-10-12 09:58
 */
public class HuffmanCode {

    public static void main(String[] args) {

//        String str = "i like like like java do you like a java";
//        byte[] bytes = str.getBytes();
//        byte[] content = huffmanZip(bytes);
//        System.out.println(Arrays.toString(content));
//
//        List<Node> nodes = getNodes(bytes);
//        Node huffmanTree = createHuffmanTree(nodes);
//        Map<Byte, String> code = getCode(huffmanTree);
//
//        byte[] decode = decode(code, content);
//        System.out.println(new String(decode));

        //测试压缩文件
//        String srcFile ="F://1.txt";
//        String dstFile = "F://dist.zip";
//        zipFile(srcFile,dstFile);
//        System.out.println("压缩文件成功！");


//        //测试解压文件
        String zipFile = "F://dist.zip";
        String dstFile ="F://2.txt";
        unZipFile(zipFile,dstFile);
        System.out.println("解压文件成功！");
    }


    /**
     * 结压文件
     *
     * @param zipFile 准备解压的文件
     * @param dstFile 将文件解压到哪个路径
     */
    public static void unZipFile(String zipFile, String dstFile) {
        //定义文件输入流
        InputStream is = null;
        //定义一个对象输入流
        ObjectInputStream ois = null;
        //定义文件输出流
        OutputStream os = null;
        try {

            //创建文件输入流
            is = new FileInputStream(zipFile);
            //创建一个和is关联的对象输入流
            ois = new ObjectInputStream(is);
            //读取byte数组 HuffmanBytes
            byte[] huffmanBytes = (byte[]) ois.readObject();
            //读取赫夫曼编码表
            Map<Byte, String> huffmanCodes = (Map<Byte, String>) ois.readObject();
            //解码
            byte[] decode = decode(huffmanCodes, huffmanBytes);
            //将decode数组写入到目标文件
            os = new FileOutputStream(dstFile);
            os.write(decode);
        } catch (Exception e) {
            System.out.println(e.getMessage());
        } finally {
            try {
                os.close();
                ois.close();
                is.close();
            } catch (IOException e) {
                System.out.println(e.getMessage());
            }

        }

    }


    /**
     * 压缩一个文件
     *缺点：1.重复数据不多的文件压缩效果不明显
     * 2.被压缩过的文件压缩效果也不明显
     * 3.可以处理所有的文件（二进制文件，文本文件）
     * @param srcFile 你传入的希望压缩的文件全路径
     * @param dstFile 我们压缩后将压缩文件放到哪个目录下
     */
    public static void zipFile(String srcFile, String dstFile) {
        FileInputStream fileInputStream = null;
        OutputStream outputStream = null;
        ObjectOutputStream objectOutputStream = null;
        try {
            //1.获取文件流
            fileInputStream = new FileInputStream(srcFile);
            //2.创建一个和文件大小的数组来存放数据
            byte[] b = new byte[fileInputStream.available()];
            //3.读取文件
            fileInputStream.read(b);
            //4.对源文件压缩
            byte[] huffmanBytes = huffmanZip(b);
            //5.创建文件输出流，存放压缩文件
            outputStream = new FileOutputStream(dstFile);
            //6.创建一个和文件输出流关联的对象输出流
            objectOutputStream = new ObjectOutputStream(outputStream);
            //7.把赫夫曼编码后的字节数组写入压缩文件
            objectOutputStream.writeObject(huffmanBytes);
            //8.以对象流的方式写入赫夫曼编码，方便后面恢复文件使用
            objectOutputStream.writeObject(huffmanCodes);
        } catch (Exception e) {
            System.out.println(e.getMessage());
        } finally {
            try {
                fileInputStream.close();
                objectOutputStream.close();
                outputStream.close();
            } catch (IOException e) {
                System.out.println(e.getMessage());
            }
        }

    }


    /**
     * 对压缩数据解码
     *
     * @param huffmanCodes 赫夫曼编码表
     * @param huffmanBytes 赫夫曼编码得到的字节数组
     * @return 原来字符串对应的数组
     */
    private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        //1.先得到huffmanBytes对应的二进制字符串
        StringBuilder stringBuilder = new StringBuilder();
        //2.将byte[]数组转成二进制字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            byte b = huffmanBytes[i];
            //判断是否是最后一个字节
            boolean flag = (i == huffmanBytes.length - 1);
            stringBuilder.append(byteToBitString(!flag, b));
        }

        //3.将赫夫曼编码表进行反转
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }

        //4. 创建集合，存放byte
        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length(); ) {
            int count = 1;//小的计数器
            boolean flag = true;
            Byte b = null;

            while (flag) {
                String key = stringBuilder.substring(i, i + count);
                b = map.get(key);
                if (b == null) {
                    count++;
                } else {
                    //匹配到了
                    flag = false;
                }
            }
            list.add(b);
            //i直接移动到count
            i += count;
        }

        //5.把list中的数据转成byte[]返回
        byte[] by = new byte[list.size()];
        for (int i = 0; i < by.length; i++) {
            by[i] = list.get(i);
        }
        return by;
    }


    /**
     * 将一个byte转换成一个二进制的字符串
     *
     * @param flag 标示是否需要补高位
     * @param b    传入的byte
     * @return 该byte对应的二进制字符串（按补码返回）
     */
    private static String byteToBitString(boolean flag, byte b) {
        //定义一个变量保存b
        int temp = b;
        //如果是正数我们需要补高位
        if (flag) {
            temp |= 256; //按位与
        }
        //返回的是temp对应的二进制的补码
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }

    /**
     * \
     * 封装
     */
    private static byte[] huffmanZip(byte[] bytes) {
        List<Node> nodes = getNodes(bytes);
        Node huffmanTree = createHuffmanTree(nodes);
        Map<Byte, String> code = getCode(huffmanTree);
        byte[] zip = zip(bytes, code);
        return zip;
    }

    /**
     * 将字符串对应的byte[] 数组，通过赫夫曼编码表，返回一个赫夫曼编码压缩后的byte[]数组
     *
     * @param bytes
     * @param huffmanCodes
     * @return
     */
    private static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
        //1.利用huffmanCodes将bytes转换成赫夫曼编码对应的字符串
        StringBuilder stringBuilder = new StringBuilder();
        //遍历byte[]数组
        for (byte b : bytes) {
            stringBuilder.append(huffmanCodes.get(b));
        }
//        System.out.println(stringBuilder.toString());

        //统计返回byte[] huffmanCodeBytes长度
        int len;
        if (stringBuilder.length() % 8 == 0) {
            len = stringBuilder.length() / 8;
        } else {
            len = stringBuilder.length() / 8 + 1;
        }
        //创建压缩后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;//记录是第几个byte
        for (int i = 0; i < stringBuilder.length(); i += 8) {
            String strByte;
            if (i + 8 > stringBuilder.length()) {
                //不够8位
                strByte = stringBuilder.substring(i);
            } else {
                strByte = stringBuilder.substring(i, i + 8);
            }
            //将strByte 转换成一个byte，放入huffmanCodeBytes中
            huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte, 2);
            index++;
        }
        return huffmanCodeBytes;
    }


    /**
     * 接受一个字节数组，转成Node数组
     *
     * @param bytes
     * @return
     */
    private static List<Node> getNodes(byte[] bytes) {
        //创建一个数组
        ArrayList<Node> nodes = new ArrayList<>();
        //遍历bytes,统计每个byte出现的次数，使用map存储
        HashMap<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if (count == null) {
                //说明map中没有存放该数据，直接将数据存入
                counts.put(b, 1);
            } else {
                counts.put(b, count + 1);
            }
        }
        //把每个键值对转成一个Node对象，并加入到nodes集合中
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }


    /**
     * 通过list构建赫夫曼树
     *
     * @param nodes
     */
    private static Node createHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            //排序，从小到大，规则是按照Node类中的排序规则
            Collections.sort(nodes);
            //取出两课最小的二叉树
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            //创建一颗新的二叉树，他的根节点没有data，只有权值
            Node parent = new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;

            //将已经处理的两颗二叉树从nodes删除
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            //将新的二叉树添加到Nodes中
            nodes.add(parent);
        }
        //返回root节点（赫夫曼树的根节点）
        return nodes.get(0);
    }


    //生成赫夫曼树对应的赫夫曼编码
    //将赫夫曼编码存放在map中
    static Map<Byte, String> huffmanCodes = new HashMap<>();
    //定义StringBuilder存储某个叶子节点的路径
    static StringBuilder stringBuilder = new StringBuilder();

    /**
     * 功能：将传入的node节点的所有叶子节点的赫夫曼编码得到，并放入到可变字节数组中
     *
     * @param node          传入的节点
     * @param code          路径： 左子节点 0  右子节点 1
     * @param stringBuilder 用来拼接路径
     */
    private static void getCodes(Node node, String code, StringBuilder stringBuilder) {
        StringBuilder stringBuilder2 = new StringBuilder(stringBuilder);
        //将code加入到stringBuilder2
        stringBuilder2.append(code);
        //如果node为空就不处理
        if (node != null) {
            //判断node为叶子结点还是非叶子节点
            if (node.data == null) {
                //是非叶子节点 进行递归处理
                //向左递归
                getCodes(node.left, "0", stringBuilder2);
                //向右递归
                getCodes(node.right, "1", stringBuilder2);
            } else {
                //说明是一个叶子节点
                //表示找到某个叶子节点的最后
                huffmanCodes.put(node.data, stringBuilder2.toString());
            }
        }
    }


    /**
     * 重载
     *
     * @param root
     * @return
     */
    private static Map<Byte, String> getCode(Node root) {
        if (root == null) {
            return null;
        }
        //向左递归
        getCodes(root.left, "0", stringBuilder);
        //向右递归
        getCodes(root.right, "1", stringBuilder);
        return huffmanCodes;
    }


    /**
     * 赫夫曼树前序遍历
     *
     * @param root
     */
    private static void preOrder(Node root) {
        //非空校验
        if (root != null) {
            root.preOrder();
        } else {
            System.out.println("赫夫曼树为空");
        }
    }

}


class Node implements Comparable<Node> {

    //存放数据本身 通过万国码表示
    Byte data;
    //权值，表示字符出现的次数
    int weight;
    Node left;
    Node right;

    public Node(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }


    /**
     * 重写比较方法
     *
     * @param o
     * @return
     */
    @Override
    public int compareTo(Node o) {
        return this.weight - o.weight;
    }


    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }


    /**
     * 前序遍历
     */
    public void preOrder() {
        System.out.println(this);

        if (this.left != null) {
            //向左递归遍历
            this.left.preOrder();
        }
        if (this.right != null) {
            //向右递归遍历
            this.right.preOrder();
        }
    }
}